#include <RcppArmadillo.h>

/* lag() takes in a matrix x with n rows and a number of lags l
 * and returns a matrix with the same dimensions as x,
 * where the first l rows are all 0 and the last n - l rows
 * are the first n - l rows of x
 */
arma::mat lag(const arma::mat& x, const int l = 1) {
    arma::mat lag_augment = arma::zeros<arma::mat>(l, x.n_cols);
    arma::mat lags = x.head_rows(x.n_rows - l);
    return arma::join_vert(lag_augment, lags);
}

/* embed() takes in a matrix x with m columns and a lag order p
 * and returns a matrix with m * p columns,
 * where the first m columns are the result of calling lag(x, 1),
 * and the last m columns are the result of calling lag(x, p)
 */
arma::mat embed(const arma::mat& x, const int p = 1) {
    const int m = x.n_cols;
    arma::mat result(x.n_rows, m*p);
    for ( arma::uword t = 0; t < p; ++t ) {
        result.cols(t*m, (t*m)+m-1) = lag(x, t+1);
    }
    return result;
}

/* lambda_tilde() takes in a matrix x of exogenous predictors,
 * a matrix beta of coefficients for those predictors,
 * (where each column is the coefficient vector for one of the equations),
 * and a matrix D giving the covariance for b,
 * and returns the matrix defined in appendix A.1 of Brandt & Sandler (2012),
 * where \( \tilde{\lambda}_{tj} = \exp(x_{tj}\beta_j) \mu_{tj}^\ast \),
 * where \( \mu^\ast = \exp( 0.5 \diag(D) ) \)
 */
arma::mat lambda_tilde(const arma::mat& x, const arma::mat& beta,
                       const arma::mat& D) {
    arma::mat log_lambda = x * beta;
    arma::rowvec b = 0.5 * D.diag().t();
    log_lambda.each_row( [b](arma::rowvec& r){ r += b; } );
    return arma::exp(log_lambda);
}

